import numpy as np
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose, BatchNormalization, Activation, UpSampling2D
from tensorflow.keras import layers
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.applications import vgg16
from tensorflow.keras import backend as K


def get_u_net(shape=(160, 112, 1), pretrained=True, freeze=True, set_trainable=None):
    # input_tensor = Input(shape=shape)
    if pretrained:
        m = vgg16.VGG16(include_top=False, weights='imagenet',
                        input_shape=shape)
    else:
        m = vgg16.VGG16(include_top=False, input_shape=shape)

    if freeze:
        for layer in m.layers:
            layer.trainable = False

    layers = dict([(layer.name, layer) for layer in m.layers])

    if not freeze and not set_trainable:
        set_trainable(layers)

    encoder_block1_conv1 = layers['block1_conv1']
    encoder_block1_conv2 = layers['block1_conv2']
    encoder_block1_pool = layers['block1_pool']

    encoder_block2_conv1 = layers['block2_conv1']
    encoder_block2_conv2 = layers['block2_conv2']
    encoder_block2_pool = layers['block2_pool']

    encoder_block3_conv1 = layers['block3_conv1']
    encoder_block3_conv2 = layers['block3_conv2']
    encoder_block3_conv3 = layers['block3_conv3']
    encoder_block3_pool = layers['block3_pool']

    encoder_block4_conv1 = layers['block4_conv1']
    encoder_block4_conv2 = layers['block4_conv2']
    encoder_block4_conv3 = layers['block4_conv3']
    encoder_block4_pool = layers['block4_pool']

    encoder_block5_conv1 = layers['block5_conv1']
    encoder_block5_conv2 = layers['block5_conv2']
    encoder_block5_conv3 = layers['block5_conv3']
    encoder_block5_pool = layers['block5_pool']

    center = Conv2D(512, (3, 3), activation='relu',
                    padding='same')(encoder_block4_pool.output)

    decoder_block4_up1 = Conv2DTranspose(
        512, (2, 2), strides=(2, 2), padding='same')(center)
    decoder_block4_concat1 = concatenate(
        [decoder_block4_up1, encoder_block4_conv3.output], axis=3)
    decoder_block4_conv1 = Conv2D(512, (3, 3), activation='relu',
                                  padding='same')(decoder_block4_concat1)
    decoder_block4_conv2 = Conv2D(512, (3, 3), activation='relu',
                                  padding='same')(decoder_block4_conv1)

    decoder_block3_up1 = Conv2DTranspose(
        256, (2, 2), strides=(2, 2), padding='same')(decoder_block4_conv2)
    decoder_block3_concat1 = concatenate(
        [decoder_block3_up1, encoder_block3_conv3.output], axis=3)
    decoder_block3_conv1 = Conv2D(256, (3, 3), activation='relu',
                                  padding='same')(decoder_block3_concat1)
    decoder_block3_conv2 = Conv2D(256, (3, 3), activation='relu',
                                  padding='same')(decoder_block3_conv1)

    decoder_block2_up1 = Conv2DTranspose(
        128, (2, 2), strides=(2, 2), padding='same')(decoder_block3_conv2)
    decoder_block2_concat1 = concatenate(
        [decoder_block2_up1, encoder_block2_conv2.output], axis=3)
    decoder_block2_conv1 = Conv2D(128, (3, 3), activation='relu',
                                  padding='same')(decoder_block2_concat1)
    decoder_block2_conv2 = Conv2D(128, (3, 3), activation='relu',
                                  padding='same')(decoder_block2_conv1)

    decoder_block1_up1 = Conv2DTranspose(
        64, (2, 2), strides=(2, 2), padding='same')(decoder_block2_conv2)
    decoder_block1_concat1 = concatenate(
        [decoder_block1_up1, encoder_block1_conv2.output], axis=3)
    decoder_block1_conv1 = Conv2D(64, (3, 3), activation='relu',
                                  padding='same')(decoder_block1_concat1)
    decoder_block1_conv2 = Conv2D(64, (3, 3), activation='relu',
                                  padding='same')(decoder_block1_conv1)

    final = Conv2D(1, (1, 1), activation='sigmoid')(decoder_block1_conv2)

    model = Model(inputs=m.input, outputs=final)

    return model


def identity_block(input_tensor, kernel_size, filters, stage, block):
    filters1, filters2 = filters
    if K.image_data_format() == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = Conv2D(filters1, kernel_size, padding='same',
                      name=conv_name_base + '2a')(input_tensor)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters2, kernel_size, padding='same',
                      name=conv_name_base + '2b')(x)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)

    x = layers.add([x, input_tensor])
    x = Activation('relu')(x)
    return x


def conv_block(input_tensor, kernel_size, filters, stage, block,
               strides=(2, 2)):
    filters1, filters2 = filters
    if K.image_data_format() == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    x = Conv2D(filters1, kernel_size, padding='same', strides=strides,
                      name=conv_name_base + '2a')(input_tensor)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
    x = Activation('relu')(x)

    x = Conv2D(filters2, kernel_size, padding='same',
                      name=conv_name_base + '2b')(x)
    x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)

    shortcut = Conv2D(filters2, (1, 1), strides=strides,
                             name=conv_name_base + '1')(input_tensor)
    shortcut = BatchNormalization(
        axis=bn_axis, name=bn_name_base + '1')(shortcut)

    x = layers.add([x, shortcut])
    x = Activation('relu')(x)
    return x


def decoder_block(input_tensor, concat_tensor, filters, name, kernel_size=3):
    x = UpSampling2D(size=(2, 2))(input_tensor)
    x = concatenate([x, concat_tensor], axis=3)
    x = Conv2D(filters, kernel_size, padding='same', name=name + '_conv1')(x)
    x = BatchNormalization(axis=3, name=name + '_bn1')(x)
    x = Activation('relu')(x)
    x = Conv2D(filters, kernel_size, padding='same', name=name + '_conv2')(x)
    x = BatchNormalization(axis=3, name=name + '_bn2')(x)
    x = Activation('relu')(x)
    return x

def res34_unet(input_shape=(224, 320, 1)):
    if K.image_data_format() == 'channels_last':
        bn_axis = 3
    else:
        bn_axis = 1

    img = Input(shape=input_shape)
    x = Conv2D(64, (7, 7), strides=(2, 2),
                      padding='same', name='conv1')(img)
    x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x)
    encoder1 = Activation('relu')(x)
    # print(encoder1.shape)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='same')(encoder1)
    # print(x.shape)
    x = conv_block(x, 3, [64, 64], stage=2, block='a', strides=(1, 1))
    x = identity_block(x, 3, [64, 64], stage=2, block='b')
    encoder2 = identity_block(x, 3, [64, 64], stage=2, block='c')

    x = conv_block(encoder2, 3, [128, 128], stage=3, block='a')
    x = identity_block(x, 3, [128, 128], stage=3, block='b')
    x = identity_block(x, 3, [128, 128], stage=3, block='c')
    encoder3 = identity_block(x, 3, [128, 128], stage=3, block='d')

    x = conv_block(encoder3, 3, [256, 256], stage=4, block='a')
    x = identity_block(x, 3, [256, 256], stage=4, block='b')
    x = identity_block(x, 3, [256, 256], stage=4, block='c')
    x = identity_block(x, 3, [256, 256], stage=4, block='d')
    x = identity_block(x, 3, [256, 256], stage=4, block='e')
    encoder4 = identity_block(x, 3, [256, 256], stage=4, block='f')

    x = conv_block(encoder4, 3, [512, 512], stage=5, block='a')
    x = identity_block(x, 3, [512, 512], stage=5, block='b')
    encoder5 = identity_block(x, 3, [512, 512], stage=5, block='c')

    decoder4 = decoder_block(encoder5, encoder4, 256, 'decoder4')
    # print(decoder4)
    decoder3 = decoder_block(decoder4, encoder3, 128, 'decoder3')
    # print(decoder3)
    decoder2 = decoder_block(decoder3, encoder2, 64, 'decoder2')
    # print(decoder2)
    decoder1 = decoder_block(decoder2, encoder1, 32, 'decoder1')
    # print(decoder1)
    
    final = Conv2D(1, (1, 1), activation='sigmoid', name='final')(decoder1)

    m = Model(img, final)

    return m



if __name__ == '__main__':
    # m = get_u_net(pretrained=False)
    m = res34_unet()
    m.summary()
